#!/usr/bin/env python3
"""
生成LangGraph工作流PNG图片
"""

import os
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.agent.langgraph_workflow import LangGraphWorkflow
from src.agent.travel_agent import TravelAgent
from src.retrieval.retrieval_tool import RetrievalTool
from src.tools.amap_navigation import AmapNavigationTool
from src.data_processing.vector_store_builder import VectorStoreBuilder
from src.models.llm_config import LLMConfig


def generate_workflow_diagram():
    """生成工作流图片"""
    try:
        print("正在初始化系统...")
        
        # 初始化系统
        config = LLMConfig()
        builder = VectorStoreBuilder(config)
        index = builder.get_index()
        retrieval_tool = RetrievalTool(index, config)
        navigation_tool = AmapNavigationTool()
        
        # 创建旅游助手
        travel_agent = TravelAgent(retrieval_tool, navigation_tool)
        
        # 创建工作流
        workflow = LangGraphWorkflow(travel_agent)
        
        print("正在生成工作流图片...")
        
        # 使用LangGraph自带的绘图功能
        graph = workflow.workflow.get_graph()
        
        # 生成Mermaid代码文件
        mermaid_path = "travel_agent_workflow.mmd"
        mermaid_code = graph.draw_mermaid()
        
        with open(mermaid_path, 'w', encoding='utf-8') as f:
            f.write(mermaid_code)
        
        print(f"✓ Mermaid代码已生成: {mermaid_path}")
        image_path = mermaid_path
        
        # 显示工作流信息
        print("\n=== 工作流节点信息 ===")
        nodes = graph.nodes
        for node_id, node_data in nodes.items():
            print(f"节点: {node_id}")
            if hasattr(node_data, 'func'):
                print(f"  函数: {node_data.func.__name__}")
            print()
        
        print("=== 工作流边信息 ===")
        edges = graph.edges
        for edge in edges:
            print(f"边: {edge}")
        
        # 显示Mermaid代码
        print("\n=== Mermaid代码 ===")
        mermaid_code = graph.draw_mermaid()
        print(mermaid_code)
        
        return image_path
        
    except Exception as e:
        print(f"生成工作流图片失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def main():
    """主函数"""
    print("=== LangGraph工作流图片生成器 ===")
    
    # 生成图片
    image_path = generate_workflow_diagram()
    
    if image_path and os.path.exists(image_path):
        print(f"\n✓ 工作流图片生成成功: {image_path}")
        print(f"文件大小: {os.path.getsize(image_path)} bytes")
    else:
        print("\n✗ 工作流图片生成失败")


if __name__ == "__main__":
    main()
